import numpy as np
import pandas as pd
import torch
from torch import nn
import torch.nn.functional as F

class FC3DDataProcessor:
    """3D彩票数据处理器 - 专门处理0-9范围的3位数字"""
    
    def __init__(self, csv_file, window_size=10):
        self.csv_file = csv_file
        self.window_size = window_size
        self.data = pd.read_csv(csv_file)
    
    def load_and_process_data(self):
        """加载和处理3D历史数据"""
        print(" 加载和处理3D历史数据...")
        
        # 检查数据格式
        required_columns = ['num_1', 'num_2', 'num_3']
        for col in required_columns:
            if col not in self.data.columns:
                raise ValueError(f"数据文件缺少 '{col}' 列")
        
        # 直接使用num_1, num_2, num_3列，提取3个数字（3D格式）
        numbers_data = []
        for _, row in self.data.iterrows():
            try:
                # 提取3个数字并确保在0-9范围内
                parsed_numbers = []
                for i in range(1, 4):  # num_1, num_2, num_3
                    col_name = f'num_{i}'
                    num = int(row[col_name])
                    # 使用模运算确保在0-9范围内
                    num = num % 10
                    parsed_numbers.append(num)
                numbers_data.append(parsed_numbers)
            except (ValueError, KeyError):
                # 如果有任何问题，添加默认值[0, 0, 0]
                numbers_data.append([0, 0, 0])
        
        numbers_array = np.array(numbers_data)
        print(f" 解析的号码数据形状: {numbers_array.shape}")
        print(f" 数据范围: {numbers_array.min()} - {numbers_array.max()}")
        
        return self.create_sequences_with_zone_transition(numbers_array)
    
    def create_sequences_with_zone_transition(self, data):
        """创建带区域转换特征的时间序列数据"""
        X, y = [], []
        
        for i in range(len(data) - self.window_size):
            # 特征：前window_size期的数据
            sequence = data[i:i + self.window_size]  # [window_size, 3]
            
            # 添加区域转换特征
            # 0-4为小数区，5-9为大数区
            zone_transitions = []
            for pos in range(3):  # 3个位置
                transitions = []
                for j in range(1, len(sequence)):
                    prev_digit = sequence[j-1][pos]
                    curr_digit = sequence[j][pos]
                    # 记录区域转换：0表示无转换，1表示小数区到大数区，-1表示大数区到小数区
                    prev_zone = 0 if prev_digit <= 4 else 1  # 0表示小数区，1表示大数区
                    curr_zone = 0 if curr_digit <= 4 else 1
                    transition = curr_zone - prev_zone  # -1, 0, 1
                    transitions.append(transition)
                zone_transitions.append(transitions)
            
            # 将区域转换特征添加到序列中
            # 扩展序列维度以包含区域转换信息
            extended_sequence = []
            for j in range(len(sequence)):
                extended_features = []
                for pos in range(3):
                    # 原始数字
                    extended_features.append(sequence[j][pos])
                    # 如果不是第一个时间步，添加区域转换特征
                    if j > 0:
                        # 获取该位置在当前时间步的区域转换
                        transition_idx = j - 1  # transitions索引
                        if transition_idx < len(zone_transitions[pos]):
                            extended_features.append(zone_transitions[pos][transition_idx])
                        else:
                            extended_features.append(0)  # 默认无转换
                    else:
                        extended_features.append(0)  # 第一个时间步无转换
                extended_sequence.append(extended_features)
            
            X.append(np.array(extended_sequence))
            
            # 标签：下一期的数据
            target = data[i + self.window_size]  # [3]
            y.append(target)
        
        X = np.array(X)  # [num_samples, window_size, 6] (3个原始数字 + 3个区域转换特征)
        y = np.array(y)  # [num_samples, 3]
        
        print(f" 序列特征形状: {X.shape}")
        print(f" 序列标签形状: {y.shape}")
        
        return X, y
    
    def get_recent_data(self):
        """获取最近的数据用于预测（使用新的区域转换特征）"""
        # 使用新的带区域转换特征的数据处理方法
        X, _ = self.load_and_process_data()
        if len(X) > 0:
            return torch.FloatTensor(X[-1:])  # 返回最后一个序列
        return None
    
    def analyze_zone_transitions(self, data):
        """分析区域转换模式"""
        transitions_stats = {
            'small_to_large': 0,  # 小数区到大数区
            'large_to_small': 0,  # 大数区到小数区
            'no_change': 0        # 无区域变化
        }
        
        for i in range(1, len(data)):
            for pos in range(3):
                prev_digit = data[i-1][pos]
                curr_digit = data[i][pos]
                
                prev_zone = 0 if prev_digit <= 4 else 1  # 0表示小数区，1表示大数区
                curr_zone = 0 if curr_digit <= 4 else 1
                
                if prev_zone == 0 and curr_zone == 1:
                    transitions_stats['small_to_large'] += 1
                elif prev_zone == 1 and curr_zone == 0:
                    transitions_stats['large_to_small'] += 1
                else:
                    transitions_stats['no_change'] += 1
        
        return transitions_stats


class FC3DSequenceLSTM(nn.Module):
    """3D彩票序列LSTM模型"""
    
    def __init__(self, input_dim=6, hidden_dim=128, num_layers=3, dropout=0.3):
        super(FC3DSequenceLSTM, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        
        # LSTM层
        self.lstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0
        )
        
        # 输出层 - 为每个位置预测数字
        self.position_heads = nn.ModuleList([
            nn.Linear(hidden_dim, 10) for _ in range(3)  # 3个位置，每个位置0-9
        ])
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        """前向传播"""
        # LSTM输出
        lstm_out, _ = self.lstm(x)  # [batch_size, seq_len, hidden_dim]
        
        # 只使用最后一个时间步的输出进行预测
        last_output = lstm_out[:, -1, :]  # [batch_size, hidden_dim]
        last_output = self.dropout(last_output)
        
        # 为每个位置生成预测
        predictions = []
        for head in self.position_heads:
            pred = head(last_output)  # [batch_size, 10]
            predictions.append(pred)
        
        # 堆叠预测结果
        predictions = torch.stack(predictions, dim=1)  # [batch_size, 3, 10]
        return predictions
    
    def predict(self, x):
        """预测函数"""
        self.eval()
        with torch.no_grad():
            logits = self.forward(x)
            probabilities = F.softmax(logits, dim=-1)
            predictions = torch.argmax(probabilities, dim=-1)
            return predictions, probabilities


def create_fc3d_sequence_model(input_dim=6, hidden_dim=128):
    """创建3D彩票序列LSTM模型"""
    model = FC3DSequenceLSTM(
        input_dim=input_dim,
        hidden_dim=hidden_dim,
        num_layers=3,
        dropout=0.3
    )
    return model


if __name__ == "__main__":
    # 测试模型创建
    print("[INFO] 测试3D彩票序列LSTM模型")
    
    # 创建模型
    model = create_fc3d_sequence_model()
    print(f"[INFO] 模型创建成功")
    print(f" 模型参数数量: {sum(p.numel() for p in model.parameters())}")
    
    # 测试前向传播 (注意：现在输入维度是6，包含3个数字特征和3个区域转换特征)
    batch_size, seq_len, input_dim = 2, 10, 6
    test_input = torch.randint(0, 10, (batch_size, seq_len, input_dim)).float()
    
    print(f" 测试输入形状: {test_input.shape}")
    
    with torch.no_grad():
        output = model(test_input)
        print(f" 模型输出形状: {output.shape}")
        
        # 测试预测
        predictions, probabilities = model.predict(test_input)
        print(f" 预测结果形状: {predictions.shape}")
        print(f" 预测示例: {predictions[0].tolist()}")